from gamma_torch import *
from opt_einsum import contract

# perams[e_sk,e_se,d_sk,d_se]





def Pi_PfPi_mul(perams_u_skse,VdVsk,VdVse):
    return contract("bcij,edij,cdx,eb->x",perams_u_skse,
                    th.conj(perams_u_skse),
                    VdVse, th.conj(VdVsk.T))
def K_PfPi_mul(perams_u_skse,perams_s_skse,VdVsk,VdVse):
    return contract("bcij,edij,cdx,eb->x",perams_u_skse,
                    th.conj(perams_s_skse),
                   VdVse, th.conj(VdVsk.T))

def Pi_PfPi(perams_u_skse,VdVsk,VdVse):
    return contract("bcij,edij,cd,eb",perams_u_skse,
                    th.conj(perams_u_skse),
                    VdVse, th.conj(VdVsk.T))
def K_PfPi(perams_u_skse,perams_s_skse,VdVsk,VdVse):
    return contract("bcij,edij,cd,eb",perams_u_skse,
                    th.conj(perams_s_skse),
                   VdVse, th.conj(VdVsk.T))

def PiKX_PxPyPz(perams_u,perams_s,VdVsk,VdVse):
    return contract("ojbc,kiba,pmda,lndc,jn,ko,mi,lp",
                                      perams_u,  th.conj(perams_u), 
                                      perams_u, th.conj(perams_s), 
                                    VdVse,VdVsk, th.conj(VdVse.T), th.conj(VdVsk.T))
# def PiKX_4P(perams_u,perams_s,VdVsk3,VdVsk4,VdVse1,VdVse2):
    # return contract("ojbc,kiba,pmda,lndc,jn,ko,mi,lp",
                                      # perams_u,  th.conj(perams_u), 
                                      # perams_u, th.conj(perams_s), 
                                    # th.conj(VdVse2.T),VdVsk3,
                                    # th.conj(VdVse1.T), VdVsk4)
def PiKA_4P(perams_u_skse,perams_s_skse,VdVsk3,VdVsk4,VdVse1,VdVse2):
    return contract("jiab,klab,nmdc,opdc,li,jk,pm,no",
                                       th.conj(perams_u_skse),   perams_u_skse,
                                      th.conj(perams_s_skse), (perams_u_skse), 
                                    VdVse1,th.conj(VdVsk3.T),
                                    VdVse2, th.conj(VdVsk4.T))
def PiKA_4P_mul(perams_u_skse,perams_s_skse,VdVsk3,VdVsk4,VdVse1,VdVse2):
    return contract("jiab,klab,nmdc,opdc,lix,jk,pmx,no->x",
                                       th.conj(perams_u_skse),   perams_u_skse,
                                      th.conj(perams_s_skse), (perams_u_skse), 
                                    VdVse1,th.conj(VdVsk3.T),
                                    VdVse2, th.conj(VdVsk4.T))
def PiKA_4P_mul_gi(perams_u_skse,perams_s_skse,VdVsk3,VdVsk4,VdVse1,VdVse2,gammas):
    return contract("jiab,klab,cfh,fr,nmgr,gdq,opdc,lix,jk,pmx,no->hqx",
                                       th.conj(perams_u_skse),   perams_u_skse,
                                       gammas,Gammas["g5"], th.conj(perams_s_skse),Gammas["g5"]@gammas,
                                       perams_u_skse, 
                                    VdVse1,th.conj(VdVsk3.T),
                                    VdVse2, th.conj(VdVsk4.T))
def PiKX_4P_mul(perams_u_skse,perams_s_skse,VdVsk3,VdVsk4,VdVse1,VdVse2):
    return contract("jiab,klac,nmdc,opdb,pix,jk,lmx,no->x",
                                       th.conj(perams_u_skse),   perams_u_skse,
                                      th.conj(perams_s_skse), (perams_u_skse), 
                                    VdVse1,th.conj(VdVsk3.T),
                                    VdVse2, th.conj(VdVsk4.T))
def PiKX_4P_mul_gi(perams_u_skse,perams_s_skse,VdVsk3,VdVsk4,VdVse1,VdVse2,gammas):
    return contract("jiab,klac,cfh,fr,nmdr,dgq,opgb,pix,jk,lmx,no->hqx",
                                       th.conj(perams_u_skse),   perams_u_skse,
                                     gammas,Gammas["g5"],  th.conj(perams_s_skse),Gammas["g5"]@gammas
                                     , (perams_u_skse), VdVse1,th.conj(VdVsk3.T),
                                    VdVse2, th.conj(VdVsk4.T))
def PiKH_4P_mul_gi(perams_u_sese,perams_u_sksk,perams_u_skse,
                perams_s_skse,VdVsk3,VdVsk4,VdVse1,VdVse2,gammas):
    return contract("jiab,afh, klfc,cgq,gr,nmdr,opdb,pix,jk,lmx,no->hqx",
                    th.conj(perams_u_sese),Gammas["g5"]@gammas,perams_s_skse,
                    gammas,Gammas["g5"],th.conj(perams_u_sksk),
                   perams_u_skse, VdVse2,th.conj(VdVsk3.T),
                                    VdVse1, th.conj(VdVsk4.T))
def PiKH_4P_mul(perams_u_sese,perams_u_sksk,perams_u_skse,
                perams_s_skse,VdVsk3,VdVsk4,VdVse1,VdVse2):
    return contract("jiab, klac,nmdc,opdb,pix,jk,lmx,no->x",
                    th.conj(perams_u_sese),perams_s_skse,th.conj(perams_u_sksk),
                   perams_u_skse, VdVse2,th.conj(VdVsk3.T),
                                    VdVse1, th.conj(VdVsk4.T))

def PiKX_4P(perams_u_skse,perams_s_skse,VdVsk3,VdVsk4,VdVse1,VdVse2):
    return contract("jiab,klac,nmdc,opdb,pi,jk,lm,no",
                                       th.conj(perams_u_skse),   perams_u_skse,
                                      th.conj(perams_s_skse), (perams_u_skse), 
                                    VdVse1,th.conj(VdVsk3.T),
                                    VdVse2, th.conj(VdVsk4.T))
def PiKH_4P(perams_u_sese,perams_u_sksk,perams_u_skse,perams_s_skse,VdVsk3,VdVsk4,VdVse1,VdVse2):
    return contract("jiab, klac,nmdc,opdb,pi,jk,lm,no",
                    th.conj(perams_u_sese),perams_s_skse,th.conj(perams_u_sksk),
                   perams_u_skse, VdVse2,th.conj(VdVsk3.T),
                                    VdVse1, th.conj(VdVsk4.T))


 # Tr[gi.Gs(ts,tk).gj.Gu(tk,ts).gk.Gu(ts,ts)]  gi=gk=g5 for pik
def PiK_K_3P(perams_u_sese,perams_u_skse,  
        perams_s_skse,VdVse1,VdVse2,VdVsk,gammas):
    return contract("ab,jkcb,cd,oide,ef,mngf,ga,mk,in,jo",
            Gammas[gammas[0]]@g5_d,th.conj(perams_s_skse),g5_d@Gammas[gammas[1]],perams_u_skse,
            Gammas[gammas[2]]@g5_d,th.conj(perams_u_sese),g5_d,VdVse1,VdVse2,th.conj(VdVsk.T))

def PiK_K_3P_mul(perams_u_sese,perams_u_skse,  
        perams_s_skse,VdVse1,VdVse2,VdVsk,gammas):
    return contract("jkca,cdh,oide,mnge,ga,mkx,inx,jo->hx",
            th.conj(perams_s_skse),g5_d@gammas,perams_u_skse,
            th.conj(perams_u_sese),g5_d,VdVse2,VdVse1,th.conj(VdVsk.T))

def PiK_K_3P_mul2(perams_u_sese,perams_u_skse,  
        perams_s_skse,VdVse1,VdVse2,VdVsk,gammas1,gammas2):
    return contract("arq, rp,jkcp,cdh,oide,mnge,ga,mkx,inx,jo->qhx",
            gammas1,g5_d,th.conj(perams_s_skse),g5_d@gammas2,perams_u_skse,
            th.conj(perams_u_sese),g5_d,VdVse2,VdVse1,th.conj(VdVsk.T))
# -Tr[g1.Gu(tk,ts).g2.Gs(ts,tk)] 
def K_g1g2_2P(perams_u_skse,perams_s_skse, VdVsk,VdVse, gammas):
    return -contract("ab,ijcb,cd,klda,ik,lj",
            Gammas[gammas[0]]@g5_d,th.conj(perams_s_skse),g5_d@Gammas[gammas[1]],perams_u_skse,
            th.conj(VdVsk.T),VdVse)

def K_g1g2_2P_mul(perams_u_skse,perams_s_skse, VdVsk,VdVse,
        gammas1,gammas2):
    return -contract("abm,bg,ijcg,cf,fhn,klha,ik,ljx->xmn",
            gammas1,g5_d,th.conj(perams_u_skse),g5_d,gammas2,perams_s_skse,
            th.conj(VdVsk.T),VdVse)

